#! /usr/bin/python3.3

# Prototype for a formulary in physics

# import python standard modules
import logging, fnmatch, re

# import sympy components (for symbolic math)
from sympy import Symbol, Eq, solve, preview, latex
from sympy.physics import units as u
from sympy.core.numbers import sympify

# import zodb
from ZODB.FileStorage import FileStorage
from ZODB.DB import DB

from persistent import Persistent
from persistent.mapping import PersistentMapping
from persistent.list import PersistentList

from transaction import commit


# import other parts of pyformulary
from utils import e2iMultiplication, Verbose  # @UnresolvedImport


class Quantity(Symbol, Verbose):
    """
    Class for physical quantities
    
    Quantity inherits from sympy.Symbol in order to allow symbolic mathematics
    such as solving an equation for a given quantity
    
    The "keyword" attribute is a tupel of string containing the usual names for this quantity.

    The fullname should be the most common (or the least confusing), as it is used as index by
    the FormulaManager. The fullname is also added to the keywords. While Keywords may occur in 
    more than one Quantity, the fullname has to be unique over the whole formulary.
    
    keyword arguments:
    value, author, source, _latex
    """
    
    __slots__ = ['name', 'fullname', 'keywords', 'value', 'is_constant', 'unit', 'author', \
                'source', 'index', 'parent', 'children', 'hierarchy']

    def __new__(cls, name, fullname, unit, *keywords, **kwargs):
        """
        explicitly call the __new__ method of the parent class because the parent class
        only accepts two arguments
        """
        assert(name and isinstance(name, str))
        assert(fullname and isinstance(fullname, str))
        
        return super(Quantity, cls).__new__(cls, name)


    def __init__(self, name, fullname, unit, *keywords, **kwargs):
        """
        Initialize a new Quantity with given values.
        For a constant specify the keyword argument "value"
        
        Fullname has to be unique for the FormulaManager
        
        The unit should be an instance of sympy.physics.units.Unit or some calculation
        based on sympy Units
        """
        
        self.name = name
        self.fullname = fullname
        self.unit = unit
        self.keywords = [fullname] + list(keywords)
        
        # hierarchy is an integer which describes the objects level. 0 for first generation,
        # 1 for child generation, 2 for grandchildren etc
        self.hierarchy = 0
        
        self.children = []
        self._addKeywordAttibutes(kwargs)
        self.is_constant = self.value != None
    
    def _addKeywordAttibutes(self, dictionary):
        """
        method to simplyfy the creation of attributes given by keyword arguments
        """
        for key, value in dictionary.items():
            setattr(self, key, value)
    
    
    def __getnewargs__(self):
        """
        this method is used to evaluate the arguments for the __new__ method
        (to deserialize objects from zodb)
        """
        return (self.name, self.fullname, self.unit)
    
    def __getstate__(self):
        """
        Method to inspect the state of an object, used to serialize the object
        """
        return {key:getattr(self, key) for key in self.__slots__}
    
    def __setstate__(self, state):
        """
        Method to set the state of an object after deserialization
        """
        for key, value in state.items():
            self.__setattr__(key, value)
    
    
    def __getattr__(self, attributeName):
        """
        This method will called if an attribute cannot be found.
        It changes the default behavior for this situation (raising an AttibuteError)
        to return None if such an attribute is specified in optionalAttributes.
        Otherwise there will still be an error.
        """
        
        # default Attribute
        if attributeName in self.__slots__:
            return None          
        else:
            raise AttributeError("Quantity instance has no attribute '%s'", attributeName)


    def createDerivedQuantity(self, fullname, index, *keywords, **kwargs):
        """
        create a derived quantity.
        index is an index in the mathematical meaning of the word, it is some letters that
        are subscript. e.g. p_h for the hydrostatic Pressure
        """
        childKeywords = (fullname,) + keywords + tuple(self.keywords)
        
        obj = self.__class__.__new__(self.__class__, self.name + '_' + 'index', fullname, self.unit)
        obj.__init__(self.name + '_' + index, fullname, self.unit,
        *childKeywords, hierarchy = self.hierarchy + 1, parent = self, index = index, **kwargs)

        #register child
        self.children.append(obj)
        
        return obj
    
    
    def getParents(self):
        """ 
        recursively get all parent objects
        """
        if self.parent == None:
            return []
        else:
            return [self.parent] + self.parent.getParents()
    
    def getChildren(self):
        """
        recursively get all child objects
        """
        offspring = [] + self.children
        for child in self.children:
            offspring += child.getChildren()
        return offspring
    
    
class Equation(Persistent, Verbose):
    """
    Class for storing physical formulas
    The actual equation is a sympy equality
    
    Physical Laws are not stored stored in that type of object, because they usually have a special name
    and cannot be solved for a certain quantity
    """

    def __init__(self, leftHandSide, rightHandSide, field, latex='', condition=''):
        
        # assertions to avoid bad arguments
        assert(field and isinstance(field, (str, tuple, list)))
        assert(isinstance(latex, str) and ('=' in latex or latex==''))
        assert(isinstance(condition, str))
        
        # if one of the sides is a string, use sympy.core.numbers.sympify to get a number
        if isinstance(leftHandSide, int):
            leftHandSide = str(leftHandSide)
        if isinstance(leftHandSide, str):
            leftHandSide = sympify(leftHandSide)
        
        if isinstance(rightHandSide, int):
            rightHandSide = str(rightHandSide)
        if isinstance(rightHandSide, str):
            rightHandSide = sympify(rightHandSide)
        
        self.equation = Eq(leftHandSide, rightHandSide)
        self.field = field
        self.condition = condition
        self.latex = latex

        # get all quantities
        self.quantities = []
        s = self.equation.atoms()
        for q in s:
            if isinstance(q, Quantity):
                self.quantities.append(q)
        
        # check if units on both sides match
        units = {q:q.unit for q in self.quantities}
        rhs_units = self.equation.rhs.subs(units)
        lhs_units = self.equation.lhs.subs(units)
        k = Symbol('k')
        unit_eq = Eq(lhs_units, (rhs_units * k) if rhs_units != 0 else k) # right hand side must not be zero
        if not solve( unit_eq, k)[0].is_Number: # units must be identical, only factor is tolerable
            raise Exception("Units of Equation don't match:" + str(lhs_units) + '!=' + str(rhs_units))
    
    def __repr__(self):
        """
        print Equation will result in 'p == F/A'
        """
        return str(self.equation)
    
    
    def __str__(self):
        """
        str(Equation) will return the equation in the form 'p == F/A'
        """
        return str(self.equation)
    
    
    def solveForQuantity(self, quantity):
        """Solves the Equation for a quantity using sympy.solve"""
        assert quantity in self.quantities
        return Equation(quantity, solve(self.equation, quantity)[0], self.field, self.condition)


    def preview(self):
        """show a preview window with a rendered latex output"""
        return preview(self.equation)

    @property
    def legend(self):
        """create a legend for the equation where all used symbols are described"""
        return [(q.name, q.fullname) for q in self.quantities]
    
    
    def _latex(self, allMultiplicationsExplicit=False):
        """
        return latex Expression for equation
        if a special latex expression is given for this equation, return the special one,
        else use the one that sympy.latex 
        
        allMultiplicationsExplicit specifies if 5*a+3*b should result in 
        "5 \\cdot a + 3 \\cdot b" (True)
        or "5 a + 3 b" (False)
        """
        if self.latex != "":
            return self.latex
        elif allMultiplicationsExplicit:
            return latex(self.equation, mul_symbol='dot', mode='inline')
        else:
            return e2iMultiplication(latex(self.equation, mul_symbol='dot', mode='inline'))


class Law(Persistent, Verbose):
    """
    This class is designed to hold values of physical laws, such as "Boyle's Law"
    
    compared to equations, there does not has to stand a quantity on both sides, 
    plus it can have a special name.
    
    """
    
    def __init__(self, name, law, quantityList, field, latex='', condition=''):
        raise NotImplementedError
    

class FormulaManager(Verbose):
    """
    Manages all Quantities, Equations and Laws
    This Object is responsible for the creation for the creation and the storing of those objects
    and provides methods to find them
    """
    
    def __init__(self, dbFile):
        # Open Database
        self._store = FileStorage(dbFile)
        self._db = DB(self._store)
        self._con = self._db.open()
        self._root = self._con.root()
        
        # Initialize Storage for Equations and Quantities
        if '_e' and '_q' not in self._root:
            self._root['_q'] = PersistentMapping()
            self._root['_e'] = PersistentList()
        
        self._q = self._root['_q']
        self._e = self._root['_e']
        
    def __enter__(self):
        '''
        method used by 'with' statement
        '''
        return self
        
    def __exit__(self, *args):
        """
        close database (used by 'with' statement)
        """
        # save changes
        commit()
        
        # close database
        self._con.close()
        self._db.close()
        self._store.close()
        
    def __getitem__(self, quantityName):
        """
        makes FormulaManager['QuantityName'] possible
        """
        return self.getQuantityByKeyword(quantityName)
        
    
    def addQuantity(self, quantity):
        """
        add a new quantity object to the internal dictionary with the fullname as key
        """
        assert isinstance(quantity, Quantity)
        assert quantity.fullname not in self._q.keys()
        self._q[quantity.fullname] = quantity


    def addEquation(self, equation):
        """
        add a new equation to the list of equations
        """
        self._e.append(equation)    
    
    
    def populateForTesting(self):
        """
        creates a (very) small set of Objects for pre-alpha testing
        """
        self.addQuantity(Quantity('p', 'Druck', u.Pa, 'Pressure'))
        self.addQuantity(Quantity('A', 'Flaeche', u.m**2))
        self.addQuantity(Quantity('F', 'Kraft', u.N))
        self.addQuantity(self.getQuantityByKeyword('Druck').createDerivedQuantity('hydrostatischer Druck', 'h'))
        self.addQuantity(Quantity('rho', 'spezifische Dichte', u.kg/u.m**3, 'Dichte'))
        self.addQuantity(Quantity('h', 'Höhe', u.m, 'Laenge', 'Distanz'))
        self.addQuantity(Quantity('g', 'Erdbeschleunigung', u.m/u.s**2,  'Fallbeschleunigung', value=9.80665))
        self.addEquation(Equation(self._q['Druck'],self._q['Kraft']/self._q['Flaeche'],'Mechanik'))
        self.addEquation(Equation(self._q['hydrostatischer Druck'], \
        self._q['spezifische Dichte']*self._q['Höhe']*self._q['Erdbeschleunigung'], 'Mechanik'))


    def getQuantitiesByKeyword(self, keyword, method='simple', number=-1, sortedBy='hierarchy'):
        """
        get all matching quantities for the given keyword
        if there are none, return empty list
        number is the maximum number of quantities returned
        
        method has to be one of ['regex', 'simple', 'wildcard']
        regex uses re
        wildcard uses fnmatch
        """
        
        if   method == 'regex'    : match = re.compile(keyword).match
        elif method == 'simple'   : match = (lambda const: lambda x: const in x)(keyword) # kind of a hack
        elif method == 'wildcard' : match = re.compile(fnmatch.translate(keyword)).match
        else                      : raise Exception('unknown Method for String Matching')
        
        r = []
        for q in self._q.values():
            for k in q.keywords:
                if match(k) and q not in r:
                    r.append(q)
            if number != -1 and len(r) >= number: # break if maximum is reached
                break
                
        r.sort(key = lambda x: getattr(x, sortedBy))
        return r


    def getQuantityByKeyword(self, keyword, method='simple'):
        """
        get a single quantity by keyword.
        Dictionary keywords (same as q.keywords[0]) trump other specified keywords.
        If multiple keywords are found, only take the first one.
        """
        if method == 'simple' and keyword in self._q: # check if keyword is in Dictionary (fast)
            return self._q[keyword]
        else:                        # loop over all quantities and check if keyword matches
            l = self.getQuantitiesByKeyword(keyword, method=method, number=-1)
            
            # avoid IndexError by checking if the list is empty and return None in that case
            if len(l) == 0:
                return None
            else:
                return l[0]


    def getEquations(self, quantity, includeChildren=True, includeParents=True):
        """
        get all equations that contain the specified quantity
        
        The argument can either be a reference to a quantity or a keyword
        """
        r = []

        if not isinstance(quantity, Quantity):
            # variable quantity is set to reference to quantity instead of quantityname
            quantity = self.getQuantityByKeyword(quantity)

        for e in self._e:
            if quantity in e.quantities:
                r.append(e)
                continue
            
            # also include equation if it contains a child quantity
            if includeChildren:
                for child in quantity.getChildren():
                    if child in e.quantities:
                        r.append(e)
                        continue
            
            # also include equation if it contains a parent quantity
            if includeParents:
                for parent in quantity.getParents(): # actually they may also be grandparents...
                    if parent in e.quantities:
                        r.append(e)
        return r
    

if __name__=="__main__":
    """
    Some playing around for testing purposes.
    Might accidentally contain old code that is not running anymore
    """

    logging.basicConfig(level=logging.INFO)

    om = FormulaManager('Database')
    om.populateForTesting()
    print("getQuantityByKeyword('Druck')", om.getQuantityByKeyword('Druck').getChildren()[0].getParents()[0])
    print("getQuantityByKeyword('*ru?k')", om.getQuantityByKeyword('*ru?k', method='wildcard'))
    print("getQuantityByKeyword('[Dd]ruc?[^x]')", om.getQuantityByKeyword('[Dd]ruc?[^x]', method='regex'))
    eqs = om.getEquations('hydrostatischer Druck')
    
    print(eqs, '\n')
    for x in eqs:
        print(x.equation)
        for q in x.legend:
            print(q[0]+":\t"+q[1])
        print('-'*30+'\n')


